"""
MCTS树节点类
表示蒙特卡洛树搜索中的一个节点
"""

import math
import random
import numpy as np
from config import WIN_COUNT, EXPLORATION_CONSTANT


class MCTSNode:
    """MCTS树节点类"""
    
    def __init__(self, board, parent=None, move=None, player=1):
        """
        初始化MCTS节点
        
        Args:
            board: 当前棋盘状态（numpy数组）
            parent: 父节点
            move: 导致此状态的落子位置 (row, col)
            player: 当前玩家（1或-1）
        """
        self.board = board.copy() if isinstance(board, np.ndarray) else board
        self.parent = parent
        self.move = move
        self.player = player
        
        self.children = []               # 子节点列表
        self.untried_moves = self.get_legal_moves()  # 未尝试的合法落子
        
        self.visits = 0                  # 访问次数
        self.wins = 0                    # 胜利次数（相对于父节点的玩家）
    
    def get_legal_moves(self):
        """
        获取所有合法落子位置（优化版）
        
        Returns:
            list: 合法落子位置列表
        """
        moves = []
        board_size = len(self.board)
        
        # 如果棋盘为空，返回中心位置
        if np.all(self.board == 0):
            center = board_size // 2
            return [(center, center)]
        
        # 只考虑已有棋子周围2格内的位置
        for i in range(board_size):
            for j in range(board_size):
                if self.board[i][j] == 0 and self.has_neighbor(i, j, distance=2):
                    moves.append((i, j))
        
        return moves if moves else []
    
    def has_neighbor(self, row, col, distance=2):
        """
        检查周围是否有棋子
        
        Args:
            row: 行坐标
            col: 列坐标
            distance: 搜索距离
        
        Returns:
            bool: 周围是否有棋子
        """
        board_size = len(self.board)
        for i in range(max(0, row - distance), 
                      min(board_size, row + distance + 1)):
            for j in range(max(0, col - distance), 
                          min(board_size, col + distance + 1)):
                if self.board[i][j] != 0:
                    return True
        return False
    
    def ucb1(self, exploration_constant=EXPLORATION_CONSTANT):
        """
        计算UCB1值（Upper Confidence Bound）
        平衡探索（exploration）与利用（exploitation）
        
        Args:
            exploration_constant: 探索常数（默认√2）
        
        Returns:
            float: UCB1值
        """
        if self.visits == 0:
            return float('inf')  # 未访问的节点优先级最高
        
        # 利用项：胜率
        exploitation = self.wins / self.visits
        
        # 探索项：鼓励访问次数少的节点
        exploration = exploration_constant * math.sqrt(
            math.log(self.parent.visits) / self.visits
        )
        
        return exploitation + exploration
    
    def select_child(self):
        """
        选择UCB1值最高的子节点
        
        Returns:
            MCTSNode: 最佳子节点
        """
        return max(self.children, key=lambda child: child.ucb1())
    
    def expand(self):
        """
        扩展一个新的子节点
        
        Returns:
            MCTSNode: 新创建的子节点
        """
        if not self.untried_moves:
            return None
        
        # 随机选择一个未尝试的落子
        move = self.untried_moves.pop(random.randint(0, len(self.untried_moves) - 1))
        
        # 创建新棋盘状态
        new_board = self.board.copy()
        new_board[move[0]][move[1]] = self.player
        
        # 创建子节点（切换玩家）
        child = MCTSNode(new_board, parent=self, move=move, player=-self.player)
        self.children.append(child)
        
        return child
    
    def is_fully_expanded(self):
        """
        检查是否所有子节点都已扩展
        
        Returns:
            bool: 是否完全扩展
        """
        return len(self.untried_moves) == 0
    
    def is_terminal(self):
        """
        检查是否为终止状态（游戏结束）
        
        Returns:
            bool: 是否为终止状态
        """
        return self.check_winner() != 0 or self.is_board_full()
    
    def check_winner(self):
        """
        检查是否有玩家获胜
        
        Returns:
            int: 1（黑方胜）、-1（白方胜）、0（未结束）
        """
        board_size = len(self.board)
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        
        for i in range(board_size):
            for j in range(board_size):
                if self.board[i][j] == 0:
                    continue
                
                player = self.board[i][j]
                
                for dx, dy in directions:
                    count = 1
                    
                    # 正方向
                    x, y = i + dx, j + dy
                    while (0 <= x < board_size and 0 <= y < board_size 
                           and self.board[x][y] == player):
                        count += 1
                        x += dx
                        y += dy
                    
                    # 反方向
                    x, y = i - dx, j - dy
                    while (0 <= x < board_size and 0 <= y < board_size 
                           and self.board[x][y] == player):
                        count += 1
                        x -= dx
                        y -= dy
                    
                    if count >= WIN_COUNT:
                        return player
        
        return 0
    
    def is_board_full(self):
        """
        检查棋盘是否已满
        
        Returns:
            bool: 棋盘是否已满
        """
        return not np.any(self.board == 0)
    
    def update(self, result):
        """
        更新节点统计信息（回溯时调用）
        
        Args:
            result: 模拟结果（1表示胜利，0表示平局，-1表示失败）
        """
        self.visits += 1
        self.wins += result
    
    def get_win_rate(self):
        """
        获取胜率
        
        Returns:
            float: 胜率（0-1之间）
        """
        if self.visits == 0:
            return 0.0
        return self.wins / self.visits
    
    def __repr__(self):
        """
        返回节点的字符串表示（用于调试）
        
        Returns:
            str: 节点信息
        """
        return f"MCTSNode(move={self.move}, visits={self.visits}, wins={self.wins}, win_rate={self.get_win_rate():.2%})"
